#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <time.h>
#include <signal.h>
#include <errno.h>
#include <pthread.h>
#include <sys/time.h>
#include <sched.h>
#include <ifaddrs.h>
#include <net/if.h>
#ifdef HAVE_GETOPT_H
#include <getopt.h>
#else
// 简单的getopt实现
extern char *optarg;
extern int optind, opterr, optopt;
int getopt(int argc, char * const argv[], const char *optstring);
#endif

#define BUFFER_SIZE 1024 * 1024  // 1MB buffer
#define DEFAULT_PORT 8888
#define MAX_CLIENTS 64

// 全局变量
volatile int running = 1;
pthread_mutex_t stats_mutex = PTHREAD_MUTEX_INITIALIZER;
long total_bytes = 0;
int active_clients = 0;
time_t start_time;

// 网络接口统计结构
typedef struct {
    char interface_name[16];
    long rx_bytes;
    long tx_bytes;
    long rx_packets;
    long tx_packets;
} interface_stats_t;

// 客户端信息结构
typedef struct {
    int client_fd;
    struct sockaddr_in client_addr;
    int client_id;
    long client_bytes;
    time_t connect_time;
    int cpu_core;  // 绑定的CPU核心
} client_info_t;

void signal_handler(int sig) {
    (void)sig;  // 避免未使用参数警告
    printf("\n正在关闭服务器...\n");
    running = 0;
}

void print_usage(const char *program_name) {
    printf("用法: %s [选项]\n", program_name);
    printf("选项:\n");
    printf("  -p <端口>    指定监听端口 (默认: %d)\n", DEFAULT_PORT);
    printf("  -h           显示此帮助信息\n");
    printf("\n注意: 此服务器支持多客户端并发连接，适合测试聚合带宽\n");
}

// 设置线程CPU亲和性
void set_thread_cpu_affinity(int cpu_core) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(cpu_core, &cpuset);
    
    if (pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset) != 0) {
        printf("警告: 无法设置线程CPU亲和性到核心%d\n", cpu_core);
    }
}

// 获取网络接口统计信息
int get_interface_stats(interface_stats_t *stats, int max_interfaces) {
    FILE *fp = fopen("/proc/net/dev", "r");
    if (!fp) return 0;
    
    char line[256];
    int count = 0;
    
    // 跳过前两行标题
    fgets(line, sizeof(line), fp);
    fgets(line, sizeof(line), fp);
    
    while (fgets(line, sizeof(line), fp) && count < max_interfaces) {
        char *interface = strtok(line, ":");
        if (!interface) continue;
        
        // 跳过回环接口
        if (strstr(interface, "lo") != NULL) continue;
        
        // 只处理物理网卡和bond接口
        if (strstr(interface, "eth") == NULL && 
            strstr(interface, "bond") == NULL &&
            strstr(interface, "ens") == NULL) continue;
        
        strcpy(stats[count].interface_name, interface);
        
        // 解析统计数据
        char *data = strtok(NULL, ":");
        if (data) {
            sscanf(data, "%ld %ld %*d %*d %*d %*d %*d %*d %ld %ld", 
                   &stats[count].rx_bytes, &stats[count].rx_packets,
                   &stats[count].tx_bytes, &stats[count].tx_packets);
            count++;
        }
    }
    
    fclose(fp);
    return count;
}

// 显示网络接口流量统计
void show_interface_stats() {
    interface_stats_t stats[16];
    int count = get_interface_stats(stats, 16);
    
    if (count == 0) return;
    
    printf("\n=== 网络接口流量统计 ===\n");
    printf("%-12s %12s %12s %12s %12s\n", 
           "接口", "接收字节", "发送字节", "接收包", "发送包");
    printf("%-12s %12s %12s %12s %12s\n", 
           "----", "--------", "--------", "------", "------");
    
    for (int i = 0; i < count; i++) {
        printf("%-12s %12ld %12ld %12ld %12ld\n",
               stats[i].interface_name,
               stats[i].rx_bytes,
               stats[i].tx_bytes,
               stats[i].rx_packets,
               stats[i].tx_packets);
    }
    printf("========================\n");
}

// 客户端处理线程
void* handle_client(void* arg) {
    client_info_t *client = (client_info_t*)arg;
    char buffer[BUFFER_SIZE];
    ssize_t bytes_received;
    time_t last_update_time = time(NULL);
    
    // 设置线程CPU亲和性
    if (client->cpu_core >= 0) {
        set_thread_cpu_affinity(client->cpu_core);
    }
    
    printf("客户端%d连接: %s:%d (CPU核心: %d)\n", 
           client->client_id,
           inet_ntoa(client->client_addr.sin_addr), 
           ntohs(client->client_addr.sin_port),
           client->cpu_core);
    
    // 更新活跃客户端计数
    pthread_mutex_lock(&stats_mutex);
    active_clients++;
    pthread_mutex_unlock(&stats_mutex);
    
    // 接收数据
    while (running) {
        bytes_received = recv(client->client_fd, buffer, BUFFER_SIZE, 0);
        if (bytes_received <= 0) {
            if (bytes_received == 0) {
                printf("客户端%d断开连接\n", client->client_id);
            } else {
                perror("接收数据失败");
            }
            break;
        }
        
        client->client_bytes += bytes_received;
        
        // 更新全局统计
        pthread_mutex_lock(&stats_mutex);
        total_bytes += bytes_received;
        pthread_mutex_unlock(&stats_mutex);
        
        // 每秒显示一次进度
        time_t current_time = time(NULL);
        if (current_time != last_update_time) {
            double duration = difftime(current_time, start_time);
            if (duration > 0) {
                pthread_mutex_lock(&stats_mutex);
                long current_total = total_bytes;
                int current_clients = active_clients;
                pthread_mutex_unlock(&stats_mutex);
                
                double bandwidth_mbps = (current_total * 8.0) / (duration * 1024 * 1024);
                printf("\r活跃客户端: %d, 总接收: %.2f MB, 聚合带宽: %.2f Mbps", 
                       current_clients, current_total / (1024.0 * 1024.0), bandwidth_mbps);
                fflush(stdout);
            }
            last_update_time = current_time;
        }
    }
    
    // 更新活跃客户端计数
    pthread_mutex_lock(&stats_mutex);
    active_clients--;
    pthread_mutex_unlock(&stats_mutex);
    
    printf("\n客户端%d统计: 接收 %.2f MB\n", 
           client->client_id, client->client_bytes / (1024.0 * 1024.0));
    
    close(client->client_fd);
    free(client);
    
    return NULL;
}

int main(int argc, char *argv[]) {
    int server_fd;
    struct sockaddr_in server_addr;
    int port = DEFAULT_PORT;
    int opt;
    int client_id = 0;
    
    // 解析命令行参数
    while ((opt = getopt(argc, argv, "p:h")) != -1) {
        switch (opt) {
            case 'p':
                port = atoi(optarg);
                if (port <= 0 || port > 65535) {
                    fprintf(stderr, "错误: 端口号必须在1-65535之间\n");
                    exit(EXIT_FAILURE);
                }
                break;
            case 'h':
                print_usage(argv[0]);
                exit(EXIT_SUCCESS);
            default:
                print_usage(argv[0]);
                exit(EXIT_FAILURE);
        }
    }
    
    // 设置信号处理
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);
    
    // 创建socket
    server_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (server_fd == -1) {
        perror("创建socket失败");
        exit(EXIT_FAILURE);
    }
    
    // 设置socket选项，允许地址重用和端口重用
    int opt_val = 1;
    if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt_val, sizeof(opt_val)) == -1) {
        perror("设置SO_REUSEADDR失败");
        close(server_fd);
        exit(EXIT_FAILURE);
    }
    
    // 设置SO_REUSEPORT，允许多个socket绑定同一端口（支持bond负载均衡）
    if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEPORT, &opt_val, sizeof(opt_val)) == -1) {
        perror("设置SO_REUSEPORT失败");
        close(server_fd);
        exit(EXIT_FAILURE);
    }
    
    // 配置服务器地址
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(port);
    
    // 绑定socket
    if (bind(server_fd, (struct sockaddr *)&server_addr, sizeof(server_addr)) == -1) {
        perror("绑定socket失败");
        close(server_fd);
        exit(EXIT_FAILURE);
    }
    
    // 开始监听
    if (listen(server_fd, MAX_CLIENTS) == -1) {
        perror("监听失败");
        close(server_fd);
        exit(EXIT_FAILURE);
    }
    
    printf("多线程TCP带宽测试服务器启动成功\n");
    printf("================================\n");
    printf("监听端口: %d\n", port);
    printf("最大并发连接: %d\n", MAX_CLIENTS);
    printf("CPU核心数: %ld\n", sysconf(_SC_NPROCESSORS_ONLN));
    printf("等待客户端连接...\n");
    printf("按Ctrl+C退出\n\n");
    
    // 显示初始网络接口状态
    show_interface_stats();
    
    start_time = time(NULL);
    
    while (running) {
        // 接受客户端连接
        struct sockaddr_in client_addr;
        socklen_t client_len = sizeof(client_addr);
        int client_fd = accept(server_fd, (struct sockaddr *)&client_addr, &client_len);
        
        if (client_fd == -1) {
            if (errno == EINTR) {
                continue;  // 被信号中断，继续循环
            }
            perror("接受连接失败");
            continue;
        }
        
        // 检查是否超过最大连接数
        pthread_mutex_lock(&stats_mutex);
        if (active_clients >= MAX_CLIENTS) {
            printf("警告: 已达到最大连接数 %d，拒绝新连接\n", MAX_CLIENTS);
            close(client_fd);
            pthread_mutex_unlock(&stats_mutex);
            continue;
        }
        pthread_mutex_unlock(&stats_mutex);
        
        // 创建客户端信息结构
        client_info_t *client = malloc(sizeof(client_info_t));
        if (!client) {
            perror("内存分配失败");
            close(client_fd);
            continue;
        }
        
        client->client_fd = client_fd;
        client->client_addr = client_addr;
        client->client_id = ++client_id;
        client->client_bytes = 0;
        client->connect_time = time(NULL);
        client->cpu_core = client_id % sysconf(_SC_NPROCESSORS_ONLN);  // 轮询分配CPU核心
        
        // 创建线程处理客户端
        pthread_t thread;
        if (pthread_create(&thread, NULL, handle_client, client) != 0) {
            perror("创建线程失败");
            free(client);
            close(client_fd);
            continue;
        }
        
        // 分离线程，让系统自动回收资源
        pthread_detach(thread);
    }
    
    close(server_fd);
    
    // 等待所有客户端断开连接
    printf("\n等待所有客户端断开连接...\n");
    while (active_clients > 0) {
        sleep(1);
    }
    
    // 显示最终统计
    time_t end_time = time(NULL);
    double duration = difftime(end_time, start_time);
    
    printf("\n=== 服务器最终统计 ===\n");
    printf("总接收数据: %.2f MB\n", total_bytes / (1024.0 * 1024.0));
    printf("运行时间: %.2f 秒\n", duration);
    if (duration > 0) {
        double bandwidth_mbps = (total_bytes * 8.0) / (duration * 1024 * 1024);
        printf("平均聚合带宽: %.2f Mbps\n", bandwidth_mbps);
        printf("平均速度: %.2f MB/s\n", total_bytes / (duration * 1024 * 1024));
    }
    printf("====================\n");
    
    // 显示最终网络接口统计
    show_interface_stats();
    
    pthread_mutex_destroy(&stats_mutex);
    printf("服务器已关闭\n");
    
    return 0;
}
